import json
import argparse
import yaml
import shutil
import sys
import os
import tempfile

# Define the argument parser
parser = argparse.ArgumentParser(description="Process a JSON file and convert it to a DataFrame.")
parser.add_argument('alias_key', type=str, help='Path to the alias_key in JSON format')
parser.add_argument('lineage_file', type=str, help='Path to the input_file containing all lineages')
parser.add_argument('--preserve-tmp-dir', action='store_true', help='Preserve the temporary directory to keep the intermediate files')

# Parse the arguments
args = parser.parse_args()

# Load the JSON data from the file path provided as argument
with open(args.alias_key, 'r') as file:
    alias_key = json.load(file)

# Create a list to store the reformatted data
reformatted_data = []

# Loop through the JSON and format it as desired
for key, value in alias_key.items():
    if value and isinstance(value, str):
        reformatted_data.append({"name": key, "alias": value})
    else:
        reformatted_data.append({"name": key, "alias": key})

# Load the lineage data from the file path provided as argument
with open(args.lineage_file, 'r') as file:
    lineages = file.read().splitlines()

# Create a dictionary to store the alias mappings
alias_dict = {item['name']: item['alias'] for item in reformatted_data}

# Function to unalias a lineage
def unalias_lineage(lineage):
    parts = lineage.split('.')
    if parts[0] in alias_dict:
        parts[0] = alias_dict[parts[0]]
    return '.'.join(parts)

# Unalias all lineages
unaliased_lineages = {lineage: unalias_lineage(lineage) for lineage in lineages}

# Function to find the immediate parent lineage
def find_immediate_parent(lineage):
    if '.' in lineage:
        return lineage.rsplit('.', 1)[0]
    return None

# Create the lineage definitions
lineage_definitions = {}
for lineage, unaliased in unaliased_lineages.items():
    parent = find_immediate_parent(unaliased)
    if parent:
        parent_unaliased = unalias_lineage(parent)
        lineage_definitions[lineage] = {"parents": [parent_unaliased]}
    else:
        lineage_definitions[lineage] = {"parents": []}

# Transform the data into YAML
yaml.dump(lineage_definitions, sys.stdout, default_flow_style=False)

if args.preserve_tmp_dir:
    temp_dir = tempfile.mkdtemp()
    print(f"Temporary directory: {temp_dir}", file=sys.stderr)
else:
    temp_dir = tempfile.mkdtemp()
    shutil.rmtree(temp_dir)